//
// Created on 2025/10/15.
//
// Node APIs are not fully supported. To solve the compilation error of the interface cannot be found,
// please include "napi/native_api.h".

#ifndef MSLITEVIDEOENHANCE_TORCH_FFT_IFFT_H
#define MSLITEVIDEOENHANCE_TORCH_FFT_IFFT_H
#include "fftw3.h"
#include <cassert>
#include <vector>
/**
 * 创建汉宁窗的平方根（与PyTorch的torch.hann_window(512).pow(0.5)一致）
 */
std::vector<float> create_sqrt_hann_window(int win_length) {
    std::vector<float> window(win_length);
    for (int i = 0; i < win_length; ++i) {
        float value = 0.5f * (1.0f - std::cos(2.0f * M_PI * i / (win_length - 1)));
        window[i] = std::sqrt(value);
    }
    return window;
}

/**
 * 执行STFT计算
 * 返回3D向量：[频率bin][时间帧][2]，其中最后一个维度是实部和虚部
 */
std::vector<std::vector<std::vector<float>>> stft(const std::vector<float> &audio_data, int sample_rate,
                                                  int n_fft = 512, int hop_length = 256, int win_length = 512,
                                                  bool center = true) {

    // 验证采样率
    assert(sample_rate == 16000 && "采样率必须为16000Hz");

    int audio_length = audio_data.size();

    // 处理center=true时的padding
    int padded_length = audio_length;
    std::vector<float> processed_audio = audio_data;

    if (center) {
        int pad_size = n_fft / 2;
        padded_length = audio_length + 2 * pad_size;
        processed_audio.resize(padded_length, 0.0f);

        // 反射padding（与PyTorch的pad_mode='reflect'一致）
        for (int i = 0; i < pad_size; ++i) {
            processed_audio[i] = audio_data[pad_size - i];                                        // 左端反射padding
            processed_audio[padded_length - 1 - i] = audio_data[audio_length - 1 - pad_size + i]; // 右端反射padding
        }
        // 将原始数据移动到中间位置
        for (int i = 0; i < audio_length; ++i) {
            processed_audio[pad_size + i] = audio_data[i];
        }
    }

    // 计算帧数
    int num_frames = 1 + (padded_length - win_length) / hop_length;
    if (num_frames < 1)
        num_frames = 1;

    // 创建窗函数
    auto window = create_sqrt_hann_window(win_length);

    // 使用单精度FFTW函数 - 关键修改
    float *fftw_input = (float *)fftwf_malloc(sizeof(float) * n_fft); // fftwf_malloc 而不是 fftw_malloc
    fftwf_complex *fftw_output = (fftwf_complex *)fftwf_malloc(sizeof(fftwf_complex) * (n_fft / 2 + 1));

    // 创建单精度FFTW计划
    fftwf_plan plan = fftwf_plan_dft_r2c_1d(n_fft, fftw_input, fftw_output, FFTW_ESTIMATE); // fftwf_plan_dft_r2c_1d
    if (!plan) {
        fftwf_free(fftw_input);  // fftwf_free
        fftwf_free(fftw_output); // fftwf_free
        throw std::runtime_error("FFTW计划创建失败");
    }

    // 初始化STFT输出结果
    int num_freq_bins = n_fft / 2 + 1;
    std::vector<std::vector<std::vector<float>>> stft_result(
        num_freq_bins, std::vector<std::vector<float>>(num_frames, std::vector<float>(2, 0.0f)));

    // 执行STFT
    for (int frame_idx = 0; frame_idx < num_frames; ++frame_idx) {
        int start_pos = frame_idx * hop_length;

        // 准备输入数据（加窗和零填充）
        for (int i = 0; i < n_fft; ++i) {
            int audio_idx = start_pos + i;
            if (i < win_length && audio_idx < padded_length) {
                fftw_input[i] = processed_audio[audio_idx] * window[i];
            } else {
                fftw_input[i] = 0.0f; // 零填充
            }
        }

        // 执行FFT - 使用单精度版本
        fftwf_execute(plan); // fftwf_execute

        // 存储结果
        for (int freq_bin = 0; freq_bin < num_freq_bins; ++freq_bin) {
            stft_result[freq_bin][frame_idx][0] = fftw_output[freq_bin][0]; // 直接使用float，无需转换
            stft_result[freq_bin][frame_idx][1] = fftw_output[freq_bin][1];
        }
    }

    // 清理FFTW资源 - 使用单精度版本
    fftwf_destroy_plan(plan); // fftwf_destroy_plan
    fftwf_free(fftw_input);   // fftwf_free
    fftwf_free(fftw_output);  // fftwf_free

    return stft_result;
}

std::vector<float> istft(const std::vector<std::vector<std::vector<float>>> &stft_result, int n_fft = 512,
                         int hop_length = 256, int win_length = 512) {

    // 验证输入维度
    if (stft_result.empty() || stft_result[0].empty()) {
        throw std::invalid_argument("STFT结果不能为空");
    }

    int num_freq_bins = stft_result.size(); // 频率bin数量 (n_fft/2 + 1)
    int num_frames = stft_result[0].size(); // 时间帧数量

    // 验证频率bin数量是否正确
    if (num_freq_bins != n_fft / 2 + 1) {
        throw std::invalid_argument("频率bin数量与n_fft不匹配");
    }

    // 创建平方根汉宁窗
    auto window = create_sqrt_hann_window(win_length);

    // 计算输出信号长度
    int output_length = (num_frames - 1) * hop_length + n_fft;

    // 初始化输出信号和窗函数累积和
    std::vector<float> output_signal(output_length, 0.0f);
    std::vector<float> window_sum(output_length, 0.0f);

    // 分配FFTW3f资源 - 关键修改：使用单精度版本
    fftwf_complex *fftw_input = (fftwf_complex *)fftwf_malloc(sizeof(fftwf_complex) * num_freq_bins);
    float *fftw_output = (float *)fftwf_malloc(sizeof(float) * n_fft);

    if (!fftw_input || !fftw_output) {
        throw std::runtime_error("FFTW内存分配失败");
    }

    // 创建IFFT计划（复数到实数变换）- 使用单精度版本
    fftwf_plan plan = fftwf_plan_dft_c2r_1d(n_fft, fftw_input, fftw_output, FFTW_ESTIMATE);
    if (!plan) {
        fftwf_free(fftw_input);
        fftwf_free(fftw_output);
        throw std::runtime_error("FFTW计划创建失败");
    }

    try {
        // 处理每一帧
        for (int frame_idx = 0; frame_idx < num_frames; ++frame_idx) {
            // 准备IFFT输入数据（复数格式）
            for (int freq_bin = 0; freq_bin < num_freq_bins; ++freq_bin) {
                // 确保有足够的数据
                if (freq_bin < stft_result.size() && frame_idx < stft_result[freq_bin].size()) {
                    fftw_input[freq_bin][0] = stft_result[freq_bin][frame_idx][0]; // 实部
                    fftw_input[freq_bin][1] = stft_result[freq_bin][frame_idx][1]; // 虚部
                } else {
                    fftw_input[freq_bin][0] = 0.0f;
                    fftw_input[freq_bin][1] = 0.0f;
                }
            }

            // 执行逆FFT - 使用单精度版本
            fftwf_execute(plan);

            // 计算当前帧在输出信号中的起始位置
            int start_pos = frame_idx * hop_length;

            // 重叠相加重建
            for (int i = 0; i < n_fft; ++i) {
                if (start_pos + i < output_length) {
                    // 应用窗函数并累加到输出信号
                    float windowed_sample = fftw_output[i] * window[i];
                    output_signal[start_pos + i] += windowed_sample;

                    // 累加窗函数平方（用于后续归一化）
                    window_sum[start_pos + i] += window[i] * window[i];
                }
            }
        }

        // 归一化处理（窗函数补偿）
        std::vector<float> reconstructed_signal(output_length);
        for (int i = 0; i < output_length; ++i) {
            if (window_sum[i] > 1e-10f) { // 避免除以零，使用单精度epsilon
                // FFTW的IFFT不包含1/N归一化，需要手动处理
                reconstructed_signal[i] = output_signal[i] / window_sum[i] / static_cast<float>(n_fft);
            } else {
                reconstructed_signal[i] = 0.0f;
            }
        }

        // 清理资源 - 使用单精度版本
        fftwf_destroy_plan(plan);
        fftwf_free(fftw_input);
        fftwf_free(fftw_output);

        return reconstructed_signal;

    } catch (...) {
        // 异常时确保资源释放
        fftwf_destroy_plan(plan);
        fftwf_free(fftw_input);
        fftwf_free(fftw_output);
        throw;
    }
}

/**
 * 将三维vector转换为float*数组
 * 内存布局: [batch][freq][time][2] -> 连续的一维数组
 *
 * @param input_data 三维vector数据，形状为 [batch_size][freq_bins][time_frames*2]
 * @param batch_size 输出参数，batch大小
 * @param freq_bins 输出参数，频率bin数量
 * @param time_frames 输出参数，时间帧数量
 * @return float* 分配的数组指针，需要调用者负责释放
 */
float *vector3d_to_float_array(const std::vector<std::vector<std::vector<float>>> &input_data, int &batch_size,
                               int &freq_bins, int &time_frames) {

    if (input_data.empty() || input_data[0].empty() || input_data[0][0].empty()) {
        throw std::invalid_argument("输入数据不能为空");
    }

    batch_size = input_data.size();
    freq_bins = input_data[0].size();
    time_frames = input_data[0][0].size() / 2; // 因为每个时间帧有实部和虚部

    // 验证数据完整性
    for (int b = 0; b < batch_size; ++b) {
        if (input_data[b].size() != freq_bins) {
            throw std::invalid_argument("频率维度不一致");
        }
        for (int f = 0; f < freq_bins; ++f) {
            if (input_data[b][f].size() != time_frames * 2) {
                throw std::invalid_argument("时间维度不一致");
            }
        }
    }

    // 分配连续内存
    int total_elements = batch_size * freq_bins * time_frames * 2;
    float *array_data = new float[total_elements];

    // 数据拷贝
    int index = 0;
    for (int b = 0; b < batch_size; ++b) {
        for (int f = 0; f < freq_bins; ++f) {
            for (int t = 0; t < time_frames * 2; ++t) {
                array_data[index++] = input_data[b][f][t];
            }
        }
    }

    return array_data;
}

/**
 * 将float*数组转换为三维vector
 * 内存布局: 连续的一维数组 -> [batch][freq][time][2]
 *
 * @param array_data float数组指针
 * @param batch_size batch大小
 * @param freq_bins 频率bin数量
 * @param time_frames 时间帧数量
 * @return std::vector<std::vector<std::vector<float>>> 三维vector
 */
std::vector<std::vector<std::vector<float>>> float_array_to_vector3d(const float *array_data, int batch_size,
                                                                     int freq_bins, int time_frames) {

    if (!array_data) {
        throw std::invalid_argument("数组指针不能为空");
    }

    if (batch_size <= 0 || freq_bins <= 0 || time_frames <= 0) {
        throw std::invalid_argument("维度参数必须为正数");
    }

    // 创建三维vector
    std::vector<std::vector<std::vector<float>>> result(
        batch_size,
        std::vector<std::vector<float>>(freq_bins, std::vector<float>(time_frames * 2, 0.0f) // 每个时间帧有实部和虚部
                                        ));

    // 数据拷贝
    int index = 0;
    for (int b = 0; b < batch_size; ++b) {
        for (int f = 0; f < freq_bins; ++f) {
            for (int t = 0; t < time_frames * 2; ++t) {
                result[b][f][t] = array_data[index++];
            }
        }
    }

    return result;
}
#endif // MSLITEVIDEOENHANCE_TORCH_FFT_IFFT_H
